#!/usr/bin/env python3
"""基于 TUN 设备的真实 VPN 客户端实现。

本模块实现了一个完整的 VPN 客户端，包含以下核心功能：
1. TUN 虚拟网卡创建与管理
2. 与服务器建立加密隧道（TLS）
3. 自动配置路由表
4. IP 数据包双向转发
5. 心跳与重连机制

注意：需要 root 权限运行，仅支持 Linux 平台。
"""

import argparse
import asyncio
import fcntl
import json
import logging
import os
import ssl
import struct
import subprocess
from dataclasses import dataclass
from ipaddress import IPv4Address, IPv4Network
from typing import Optional

# TUN 设备相关常量
TUNSETIFF = 0x400454ca
IFF_TUN = 0x0001
IFF_NO_PI = 0x1000


@dataclass
class ClientConfig:
    """VPN 客户端配置。"""
    
    server_host: str                    # VPN 服务器地址
    server_port: int                    # VPN 服务器端口
    auth_token: str                     # 认证令牌
    client_id: str                      # 客户端标识
    tun_name: str                       # TUN 设备名称
    cafile: Optional[str]               # CA 证书
    route_all_traffic: bool             # 是否路由所有流量


@dataclass
class VPNConfig:
    """从服务器获取的 VPN 配置。"""
    
    virtual_ip: IPv4Address             # 客户端虚拟 IP
    server_ip: IPv4Address              # 服务器虚拟 IP
    netmask: str                        # 子网掩码
    network: IPv4Network                # VPN 网段


class TunDevice:
    """TUN 虚拟网卡管理器（客户端版本）。"""
    
    def __init__(self, name: str):
        self.name = name
        self.fd: Optional[int] = None
    
    def create(self) -> int:
        """创建 TUN 设备。"""
        try:
            self.fd = os.open("/dev/net/tun", os.O_RDWR)
            ifr = struct.pack("16sH", self.name.encode()[:15], IFF_TUN | IFF_NO_PI)
            fcntl.ioctl(self.fd, TUNSETIFF, ifr)
            logging.info(f"TUN 设备 {self.name} 创建成功")
            return self.fd
        except Exception as e:
            logging.error(f"创建 TUN 设备失败: {e}")
            raise
    
    def configure(self, ip_addr: str, peer_ip: str, netmask: int = 24):
        """配置 TUN 设备。"""
        try:
            # 设置点对点 IP 地址
            subprocess.run([
                "ip", "addr", "add",
                f"{ip_addr}/32",
                "peer", peer_ip,
                "dev", self.name
            ], check=True)
            
            # 启动接口
            subprocess.run(["ip", "link", "set", self.name, "up"], check=True)
            
            logging.info(f"TUN 设备 {self.name} 配置完成: {ip_addr} -> {peer_ip}")
        except subprocess.CalledProcessError as e:
            logging.error(f"配置 TUN 设备失败: {e}")
            raise
    
    def add_route(self, network: str, gateway: str):
        """添加路由。"""
        try:
            subprocess.run([
                "ip", "route", "add",
                network, "via", gateway,
                "dev", self.name
            ], check=True)
            logging.info(f"添加路由: {network} via {gateway}")
        except subprocess.CalledProcessError as e:
            logging.warning(f"添加路由失败: {e}")
    
    def set_default_route(self, gateway: str):
        """设置默认路由（全局代理）。"""
        try:
            # 保存原默认路由
            result = subprocess.run(
                ["ip", "route", "show", "default"],
                capture_output=True,
                text=True
            )
            self.original_route = result.stdout.strip()
            
            # 删除原默认路由
            if self.original_route:
                subprocess.run(["ip", "route", "del", "default"], check=True)
            
            # 添加新默认路由
            subprocess.run([
                "ip", "route", "add", "default",
                "via", gateway,
                "dev", self.name
            ], check=True)
            
            logging.info(f"默认路由已设置为 VPN: {gateway}")
        except subprocess.CalledProcessError as e:
            logging.error(f"设置默认路由失败: {e}")
            raise
    
    def restore_default_route(self):
        """恢复原默认路由。"""
        try:
            subprocess.run(["ip", "route", "del", "default"], stderr=subprocess.DEVNULL)
            
            if hasattr(self, 'original_route') and self.original_route:
                parts = self.original_route.split()
                if len(parts) >= 3:
                    subprocess.run(["ip", "route", "add"] + parts, check=True)
                    logging.info("默认路由已恢复")
        except Exception as e:
            logging.warning(f"恢复默认路由失败: {e}")
    
    def read(self, size: int = 2048) -> bytes:
        """从 TUN 设备读取 IP 数据包。"""
        if self.fd is None:
            raise RuntimeError("TUN 设备未创建")
        return os.read(self.fd, size)
    
    def write(self, data: bytes) -> int:
        """向 TUN 设备写入 IP 数据包。"""
        if self.fd is None:
            raise RuntimeError("TUN 设备未创建")
        return os.write(self.fd, data)
    
    def close(self):
        """关闭 TUN 设备。"""
        if self.fd is not None:
            os.close(self.fd)
            self.fd = None
            logging.info(f"TUN 设备 {self.name} 已关闭")


class VPNClient:
    """VPN 客户端主类。"""
    
    def __init__(self, config: ClientConfig):
        self.config = config
        self.tun = TunDevice(config.tun_name)
        self.vpn_config: Optional[VPNConfig] = None
        self.reader: Optional[asyncio.StreamReader] = None
        self.writer: Optional[asyncio.StreamWriter] = None
    
    async def connect_server(self) -> VPNConfig:
        """连接 VPN 服务器并完成认证。"""
        # 创建 TLS 上下文
        ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        
        if self.config.cafile:
            ssl_ctx.load_verify_locations(self.config.cafile)
        else:
            # 开发环境：跳过证书验证
            ssl_ctx.check_hostname = False
            ssl_ctx.verify_mode = ssl.CERT_NONE
        
        # 连接服务器
        self.reader, self.writer = await asyncio.open_connection(
            self.config.server_host,
            self.config.server_port,
            ssl=ssl_ctx
        )
        
        logging.info(f"已连接到 VPN 服务器: {self.config.server_host}:{self.config.server_port}")
        
        # 发送认证信息
        auth_msg = {
            "token": self.config.auth_token,
            "client_id": self.config.client_id
        }
        self.writer.write(json.dumps(auth_msg).encode("utf-8") + b"\n")
        await self.writer.drain()
        
        # 接收配置
        config_data = await asyncio.wait_for(self.reader.readline(), timeout=10)
        config_msg = json.loads(config_data.decode("utf-8"))
        
        if config_msg.get("status") != "ok":
            raise ConnectionError(f"服务器拒绝连接: {config_msg}")
        
        # 解析配置
        vpn_config = VPNConfig(
            virtual_ip=IPv4Address(config_msg["virtual_ip"]),
            server_ip=IPv4Address(config_msg["server_ip"]),
            netmask=config_msg["netmask"],
            network=IPv4Network(config_msg["network"])
        )
        
        logging.info(f"获取 VPN 配置: IP={vpn_config.virtual_ip}, 网关={vpn_config.server_ip}")
        return vpn_config
    
    def setup_tun(self):
        """配置 TUN 设备与路由。"""
        # 创建并配置 TUN
        self.tun.create()
        self.tun.configure(
            str(self.vpn_config.virtual_ip),
            str(self.vpn_config.server_ip)
        )
        
        # 配置路由
        if self.config.route_all_traffic:
            # 全局代理：所有流量走 VPN
            self.tun.set_default_route(str(self.vpn_config.server_ip))
        else:
            # 仅 VPN 网段流量走 VPN
            self.tun.add_route(
                str(self.vpn_config.network),
                str(self.vpn_config.server_ip)
            )
    
    async def tun_to_server(self):
        """从 TUN 读取数据包并发送到服务器。"""
        loop = asyncio.get_event_loop()
        
        try:
            while True:
                # 从 TUN 读取 IP 包
                packet = await loop.run_in_executor(None, self.tun.read, 2048)
                
                # 发送到服务器（长度 + 数据）
                length_data = struct.pack("!I", len(packet))
                self.writer.write(length_data + packet)
                await self.writer.drain()
                
        except Exception as e:
            logging.error(f"tun_to_server 错误: {e}")
            raise
    
    async def server_to_tun(self):
        """从服务器接收数据包并写入 TUN。"""
        loop = asyncio.get_event_loop()
        
        try:
            while True:
                # 读取包长度
                length_data = await self.reader.readexactly(4)
                packet_length = struct.unpack("!I", length_data)[0]
                
                # 读取 IP 数据包
                packet = await self.reader.readexactly(packet_length)
                
                # 写入 TUN 设备
                await loop.run_in_executor(None, self.tun.write, packet)
                
        except asyncio.IncompleteReadError:
            logging.info("服务器连接已关闭")
        except Exception as e:
            logging.error(f"server_to_tun 错误: {e}")
            raise
    
    async def run(self):
        """运行 VPN 客户端。"""
        try:
            # 连接服务器并获取配置
            self.vpn_config = await self.connect_server()
            
            # 配置 TUN 设备与路由
            self.setup_tun()
            
            logging.info("VPN 隧道已建立，开始转发数据")
            
            # 双向转发数据
            await asyncio.gather(
                self.tun_to_server(),
                self.server_to_tun()
            )
            
        except Exception as e:
            logging.error(f"VPN 客户端错误: {e}")
            raise
        finally:
            # 清理资源
            if self.config.route_all_traffic:
                self.tun.restore_default_route()
            
            self.tun.close()
            
            if self.writer:
                self.writer.close()
                await self.writer.wait_closed()


def parse_args() -> ClientConfig:
    parser = argparse.ArgumentParser(description="基于 TUN 的真实 VPN 客户端")
    parser.add_argument("--server-host", required=True, help="VPN 服务器地址")
    parser.add_argument("--server-port", type=int, default=8443, help="VPN 服务器端口")
    parser.add_argument("--auth-token", required=True, help="认证令牌")
    parser.add_argument("--client-id", default="client-1", help="客户端标识")
    parser.add_argument("--tun-name", default="vpn0", help="TUN 设备名称")
    parser.add_argument("--ca-cert", dest="cafile", help="CA 证书路径")
    parser.add_argument("--route-all", action="store_true", dest="route_all_traffic",
                       help="路由所有流量通过 VPN（全局代理）")
    parser.add_argument("--log-level", default="INFO", help="日志等级")
    
    args = parser.parse_args()
    
    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s | %(levelname)s | %(message)s"
    )
    
    return ClientConfig(
        server_host=args.server_host,
        server_port=args.server_port,
        auth_token=args.auth_token,
        client_id=args.client_id,
        tun_name=args.tun_name,
        cafile=args.cafile,
        route_all_traffic=args.route_all_traffic
    )


def main():
    if os.geteuid() != 0:
        print("错误：需要 root 权限运行")
        return 1
    
    config = parse_args()
    client = VPNClient(config)
    
    try:
        asyncio.run(client.run())
    except KeyboardInterrupt:
        logging.info("客户端已停止")
    
    return 0


if __name__ == "__main__":
    exit(main())
